// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <cuda_runtime.h> // @manual
#include <pollux/testing/gtest_utils.h>

#include <pollux/exec/exchange_source.h>
#include <pollux/exec/plan_node_stats.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/testing/exec/util/hive_connector_test_base.h>
#include <pollux/testing/exec/util/local_exchange_source.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/experimental/wave/exec/ToWave.h>
#include <pollux/experimental/wave/exec/WaveHiveDataSource.h>
#include <pollux/experimental/wave/exec/tests/utils/FileFormat.h>
#include <pollux/experimental/wave/exec/tests/utils/WaveTestSplitReader.h>
#include <pollux/vector/vector_fuzzer.h>

TURBO_DECLARE_FLAG(int32_t, wave_max_reader_batch_rows);
TURBO_DECLARE_FLAG(int32_t, max_streams_per_driver);
TURBO_DECLARE_FLAG(int32_t, wave_reader_rows_per_tb);

using namespace kumo::pollux;
using namespace kumo::pollux::core;
using namespace kumo::pollux::exec;
using namespace kumo::pollux::wave;

class BarrierTest : public testing::Test {
 protected:
  void testFunc(int32_t threadIdx) {
    auto barrierIdx = threadIdx / 10;
    auto name = fmt::format("bar{}", barrierIdx);
    auto barrier = WaveBarrier::get(name, 0, 0);
    message(threadIdx, "enter");
    barrier->enter();
    for (auto action = 0; action < 100; ++action) {
      // std::this_thread::sleep_for(std::chrono::milliseconds(1)); // NOLINT
      if (action % 10 == 0 && threadIdx % 10 < 5) {
        void* reason = reinterpret_cast<void*>(action + 1);
        message(threadIdx, "acq");
        barrier->acquire(nullptr, reason, nullptr);
        message(threadIdx, "exc");
        // std::this_thread::sleep_for(std::chrono::milliseconds(10)); // NOLINT
        ++numAcquired_;
        message(threadIdx, "rel");
        barrier->release();
      } else if (action % 11 == 0) {
        message(threadIdx, "leave-middle");
        barrier->leave();
        message(threadIdx, "enter-middle");
        barrier->enter();
      } else {
        message(threadIdx, "mayYield");
        barrier->mayYield(nullptr, nullptr);
        message(threadIdx, "cont");
      }
    }
    message(threadIdx, "leave");
    barrier->leave();
  }

  void message(int threadIdx, const char* s) {
    if (!verbose_) {
      return;
    }
    std::lock_guard<std::mutex> l(mtx_);
    std::cout << s << " " << threadIdx << std::endl;
  }

  bool verbose_{false};
  std::atomic<int32_t> numAcquired_{0};
  std::mutex mtx_;
};

TEST_F(BarrierTest, basic) {
  constexpr int32_t kNumThreads = 100;
  std::vector<std::thread> threads;
  threads.reserve(kNumThreads);
  for (int32_t threadIndex = 0; threadIndex < kNumThreads; ++threadIndex) {
    threads.push_back(
        std::thread([this, threadIndex]() { testFunc(threadIndex); }));
  }
  for (auto& thr : threads) {
    thr.join();
  }
  EXPECT_EQ(5 * kNumThreads, numAcquired_);
}
